﻿using System;
using System.Collections.Generic;
using System.Text;
using SlimDX;
using SlimDX.Direct3D9;
using System.Runtime.InteropServices;
using System.Drawing;
using System.Windows.Forms;

namespace SlimWater
{
    struct WaterVertex
    {
        public Vector3 Position;
        public Vector3 Normal;
        public Vector2 TexCoord0;

        public WaterVertex(Vector3 Position, Vector3 Normal, Vector2 TexCoord0)
        {
            this.Position = Position;
            this.Normal = Normal;
            this.TexCoord0 = TexCoord0;
        }
    }

    class Water : IDisposable
    {
        Device device;

        VertexDeclaration vertexDecl;
        Effect effect;

        VertexBuffer vb;
        IndexBuffer ib;

        Surface refractionTarget, reflectionTarget, screen;
        public Texture refractionTex, reflectionTex, bumpMap;
        Material material;

        Matrix originalView;
        Matrix reflectView;

        Matrix world;

        public Water(Device device, float scale, int windowWidth, int windowHeight)
        {
            this.device = device;

            vb = new VertexBuffer(device, 4 * Marshal.SizeOf(typeof(WaterVertex)),
                Usage.WriteOnly, VertexFormat.None, Pool.Managed);
            vb.DebugName = "Water VertexBuffer";

            Vector3 up = new Vector3(0, 1, 0);
            var stream = vb.Lock(0, 0, LockFlags.None);
            stream.WriteRange(new[] {
				new WaterVertex(new Vector3(-1, 0, -1), up, new Vector2(0, 0)),
                new WaterVertex(new Vector3(-1, 0, 1), up, new Vector2(0, 1)),
                new WaterVertex(new Vector3(1, 0, -1), up, new Vector2(1, 0)),
                new WaterVertex(new Vector3(1, 0, 1), up, new Vector2(1, 1)),
			});
            vb.Unlock();

            ib = new IndexBuffer(device, 6 * 4, Usage.None, Pool.Default, false);
            ib.DebugName = "Water IndexBuffer";
            stream = ib.Lock(0, 0, LockFlags.None);
            stream.WriteRange(new[] {
                0, 1, 2, 1, 3, 2
            });
            ib.Unlock();

            vertexDecl = new VertexDeclaration(device, new[] {
				new VertexElement(0, 0, DeclarationType.Float3, DeclarationMethod.Default, DeclarationUsage.Position, 0), 
                new VertexElement(0, 12, DeclarationType.Float3, DeclarationMethod.Default, DeclarationUsage.Normal, 0), 
				new VertexElement(0, 24, DeclarationType.Float2, DeclarationMethod.Default, DeclarationUsage.TextureCoordinate, 0), 
                new VertexElement(0, 32, DeclarationType.Float2, DeclarationMethod.Default, DeclarationUsage.TextureCoordinate, 1), 
				VertexElement.VertexDeclarationEnd
        	});
            
            bumpMap = Texture.FromFile(device, "../../data/waterbump.dds");
            bumpMap.DebugName = "Water BumpMap Texture";

            material = new Material();
            material.Ambient = Color.Black;
            material.Diffuse = Color.Blue;

            screen = device.GetRenderTarget(0);

            refractionTex = new Texture(device, windowWidth, windowHeight, 1, Usage.RenderTarget, Format.A8R8G8B8, Pool.Default);
            reflectionTex = new Texture(device, windowWidth, windowHeight, 1, Usage.RenderTarget, Format.A8R8G8B8, Pool.Default);
            refractionTex.DebugName = "Water Refraction Texture";
            reflectionTex.DebugName = "Water Reflection Texture";

            refractionTarget = refractionTex.GetSurfaceLevel(0);
            reflectionTarget = reflectionTex.GetSurfaceLevel(0);
            refractionTarget.DebugName = "Refraction Target";
            reflectionTarget.DebugName = "Reflection Target";

            effect = Effect.FromFile(device, "../../data/water.fx", ShaderFlags.None);

            world = Matrix.Scaling(scale, scale, scale);// *Matrix.Translation(256, 0, 256);
        }

        public void Render(Vector3 cameraEP, double time)
        {
            Matrix view = device.GetTransform(TransformState.View);
            Matrix projection = device.GetTransform(TransformState.Projection);
            Matrix wvp = world * view * projection;
            Matrix mReflectMatrix = world * reflectView * projection;

            device.SetStreamSource(0, vb, 0, Marshal.SizeOf(typeof(WaterVertex)));
            device.Indices = ib;
            device.VertexDeclaration = vertexDecl;

            effect.Technique = "WaterTechnique";
            effect.SetValue("xWorldViewProj", wvp);
            effect.SetValue("xReflectMatrix", mReflectMatrix);

            effect.SetValue("xCamPosition", new Vector4(cameraEP, 1f));
            effect.SetValue("xTime", (float)time); // cast is necessary
            
            // materials & light
            Light light = device.GetLight(0);
            effect.SetValue("xMaterialAmbient", material.Ambient);
            effect.SetValue("xMaterialDiffuse", material.Diffuse);
            effect.SetValue("xMaterialSpecular", material.Specular);
            effect.SetValue("xLight", new Vector4(Vector3.Normalize(light.Direction), 1.0f));
            effect.SetValue("xLightPos", new Vector4(light.Position, 1.0f));
            effect.SetValue("xLightAmbient", light.Ambient);
            effect.SetValue("xLightDiffuse", light.Diffuse);
            effect.SetValue("xLightSpecular", light.Specular);

            // reflection & refraction textures
            effect.SetTexture("reflectTex", reflectionTex);
            effect.SetTexture("refractTex", refractionTex);
            effect.SetTexture("waterBumpTex", bumpMap);

            effect.Begin(0);
            effect.BeginPass(0);
            {
                device.DrawIndexedPrimitives(PrimitiveType.TriangleList, 0, 0, 4, 0, 2);
            }
            effect.EndPass();
            effect.End();
        }

        private Plane CalculatePlane(float height, int side, Matrix matView)
        {
            Vector4 parameters = new Vector4(0, side, 0, height);

            Matrix mViewProj = device.GetTransform(TransformState.World) * matView * device.GetTransform(TransformState.Projection);
            Matrix mInvTransposedViewProj = Matrix.Transpose(Matrix.Invert(mViewProj));

            return new Plane(Vector4.Transform(parameters, mInvTransposedViewProj));
        }

        public void RenderRefractionBegin(Vector3 cameraEP)
        {
            Plane refract;
            if (cameraEP.Y > 0) refract = CalculatePlane(3.0f, -1, device.GetTransform(TransformState.View));
            else refract = CalculatePlane(1.0f, 1, device.GetTransform(TransformState.View));

            device.SetClipPlane(0, refract);
            device.SetRenderState(RenderState.ClipPlaneEnable, 1);
            device.SetRenderTarget(0, refractionTarget);

            device.Clear(ClearFlags.Target | ClearFlags.ZBuffer, Color.Black, 1.0f, 0);
            device.BeginScene();
        }

        public void RenderRefractionEnd()
        {
            device.EndScene();
            device.SetRenderState(RenderState.ClipPlaneEnable, 0);
            device.SetRenderTarget(0, screen);
        }

        public void RenderReflectionBegin(Vector3 cameraEP, Vector3 cameraLAP)
        {
            Vector3 reflEP = cameraEP;
            reflEP.Y = -reflEP.Y;
            Vector3 reflLAP = cameraLAP;
            reflLAP.Y = -reflLAP.Y;

            originalView = device.GetTransform(TransformState.View); // backup

            reflectView = Matrix.LookAtLH(reflEP, reflLAP, new Vector3(0, 1, 0));
            device.SetTransform(TransformState.View, reflectView);

            Plane reflect = CalculatePlane(0.5f, 1, reflectView);
            device.SetClipPlane(0, reflect);
            device.SetRenderState(RenderState.ClipPlaneEnable, 1);
            device.SetRenderTarget(0, reflectionTarget);

            device.Clear(ClearFlags.Target | ClearFlags.ZBuffer, Color.Black, 1.0f, 0);
            device.BeginScene();
        }

        public void RenderReflectionEnd()
        {
            device.EndScene();
            device.SetRenderState(RenderState.ClipPlaneEnable, 0);
            device.SetRenderTarget(0, screen);
            device.SetTransform(TransformState.View, originalView); // revert back to original
        }

        public void Dispose()
        {
            effect.Dispose();

            vb.Dispose();
            ib.Dispose();
            vertexDecl.Dispose();

            reflectionTex.Dispose();
            refractionTex.Dispose();
            bumpMap.Dispose();

            refractionTarget.Dispose();
            reflectionTarget.Dispose();
            screen.Dispose();
        }
    }
}
